import os
import re
import cv2
import easyocr
import numpy as np
import pytesseract
from matplotlib import pyplot as plt


def show_img(img):
    # 检查图像通道数
    if len(img.shape) == 2:  # 单通道灰度图
        plt.imshow(img, cmap='gray')
    else:  # 彩色图
        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))  # 转换BGR到RGB
    plt.axis('off')
    plt.show()


def remove_bar(bar_area, bar_mask):
    """除去原图中bar部分，防止对后期产生干扰"""
    gray = bar_area.flatten()
    mask = bar_mask.flatten()
    valid = gray[mask == 0]
    mean = float(valid.mean())
    print(mean)
    result = bar_area.copy()
    result[bar_mask != 0] = mean
    return result


def detect_scale_bar(gray_img, img_type, specify_nm=0.0):
    """比例尺检测"""
    diff_type = {
        # 不同图片类型
        2048: {
            # 比例尺区域
            'bar': {
                'x': 115,
                'y': 1800,
                'w': 640,
                'h': 121,
            },
            'parameter': {
                'bar_max': 0  # 比例尺bar颜色最大值
            }
        },
        302: {
            'bar': {
                'x': 7,
                'y': 264,
                'w': 87,
                'h': 20,
            },
            'parameter': {
                'bar_max': 40
            },
        },
    }
    bar_x = diff_type[img_type]['bar']['x']
    bar_y = diff_type[img_type]['bar']['y']
    bar_w = diff_type[img_type]['bar']['w']
    bar_h = diff_type[img_type]['bar']['h']
    bar_max = diff_type[img_type]['parameter']['bar_max']
    region = gray_img[bar_y:bar_y + bar_h, bar_x:bar_x + bar_w]

    dark_mask = (region <= bar_max).astype(np.uint8)
    # show_img(dark_mask)

    cleaned_region = remove_bar(region, dark_mask)
    gray_img[bar_y:bar_y + bar_h, bar_x:bar_x + bar_w] = cleaned_region
    # show_img(gray_img)

    contours, _ = cv2.findContours(dark_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    max_width = 0
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if w > max_width:
            max_width = w

    if not specify_nm:
        # 没有指明比例尺实际nm长度
        # OCR文字识别
        reader = easyocr.Reader(['en'])

        ocr_text_list = reader.readtext(region, detail=0, allowlist='0125.nm')
        ocr_text = ''.join(ocr_text_list)
        print(f'ORC识别内容:{ocr_text}')
        match = re.search(r"(\d+(\.\d+)?)\s*(nm)?", ocr_text.lower())
        if match:
            scale_nm_raw = float(match.group(1))
            allowed = [5, 10, 20, 50, 100, 200]
            scale_nm = min(allowed, key=lambda x: abs(x - scale_nm_raw))
            print(f"提取数值: {scale_nm_raw} → 匹配最接近合法值: {scale_nm} nm")
        else:
            print(f"OCR识别失败")
            scale_nm = 0.0
    else:
        print(f"使用提供值: {specify_nm} nm")
        scale_nm = specify_nm

    if max_width == 0:
        raise ValueError("比例尺检测失败：未找到明显黑条。")
    print(f"比例尺条像素长度: {max_width} px")
    return gray_img,scale_nm, max_width